#!/usr/bin/python
# -*- coding: utf-8 -*-

import time
import queue
import contextlib
import threading


class Task(object):
    def __init__(self, pool, func, args=(), kwargs={}):
        '''
        初始化任务
        :param func: 函数名称
        :param args: 函数参数
        :param kwargs: 函数参数
        '''
        super().__init__()

        # 线程池状态
        self.READY = 0
        self.RUN = 1
        self.SUCCESS = 2
        self.EXCEPTION = 3
        self.TERMINATE = 4
        self._state = self.READY

        self._event = threading.Event()
        self._pool = pool
        self._func = func
        self._args = args
        self._kwargs = kwargs
        self._success = None
        self._result = None
        self._exception = None

    def exec(self):
        if self.READY != self._state:
            raise ValueError("Task not ready.")
        try:
            self._state = self.RUN
            self._result = self._func(*self._args, **self._kwargs)
            self._state = self.SUCCESS
            self._exception = None
            self._success = True
        except Exception as e:
            self._state = self.EXCEPTION
            self._exception = e
            self._result = None
            self._success = False
        finally:
            self._event.set()

    def get_state(self):
        return self._state

    def is_ready(self):
        return self.READY == self._state

    def terminate(self):
        self._state = self.TERMINATE
        self._result = None
        self._success = False
        self._event.set()

    def get_success(self):
        self._event.wait()
        return self._success

    def get_result(self):
        self._event.wait()
        return self._result

    def get_exception(self):
        self._event.wait()
        return self._exception


# 参考 http://www.cnblogs.com/hjc4025/p/6950157.html
class ThreadPool(object):
    def __init__(self, pool_size, task_queue_size=0) -> None:
        '''
        初始化线程池
        :param pool_size: 线程池的大小, 即线程池中可容纳线程的最大数量
        :param task_queue_size: 任务队列的大小, 即任务队列中可以保存的最大数量
        '''
        super().__init__()

        # 线程池状态
        self.RUN = 0
        self.CLOSE = 1
        self.TERMINATE = 2

        # 停止任务
        self.STOP_TASK = object()

        # 设置线程池的大小
        self._pool_size = pool_size
        # 创建任务队列
        self._task_queue = queue.Queue(task_queue_size)
        self._all_task_list = []
        # 线程池状态
        self._state = self.RUN
        # 空闲线程列表
        self._free_thread_list = []
        # 所有线程列表
        self._all_thread_list = []
        # event, 默认为不用等待, 当有线程启动后, 调用clear()方法
        self._event = threading.Event()
        self._event.set()

    def apply(self, func, args=(), kwargs={}):
        '''
        向队列中添加一个任务, 等待任务完成, 返回任务的返回值
        :param func: 
        :param args: 
        :param kwargs: 
        :return: 
        '''
        return self.apply_async(func, args, kwargs).get_result()

    def apply_async(self, func, args=(), kwargs={}):
        '''
        向队列中添加一个任务, 立即返回
        :param func: 函数名称
        :param args: 函数参数
        :param kwargs: 函数参数
        '''
        if self.RUN != self._state:
            raise ValueError("thread pool not running.")

        # 如果没有空闲的线程, 并且现有的线程数量小于最大线程数，则创建新线程
        if 0 == len(self._free_thread_list) and len(self._all_thread_list) < self._pool_size:
            self._generate_thread()

        # 构造任务
        task = Task(self, func, args, kwargs)
        # 将任务放入任务队列
        self._task_queue.put(task)
        self._all_task_list.append(task)
        return task

    def wait(self):
        self._event.wait()

    def close(self):
        '''
        关闭线程池, 会等待所有任务完成
        '''
        self._state = self.CLOSE
        for i in range(len(self._all_thread_list)):
            self._task_queue.put(self.STOP_TASK)

    def terminate(self):
        '''
        强制终止线程池
        该函数保证调用结束后, 线程池中的线程全部都立即终止掉
        '''
        self._state = self.TERMINATE
        # 等待线程终止
        self.wait()

        for task in self._all_task_list:
            if isinstance(task, Task):
                if task.is_ready():
                    task.terminate()

    def task_list(self):
        '''
        获得任务列表
        :return: 任务列表
        '''
        return self._all_task_list

    def _generate_thread(self):
        '''
        创建线程
        '''
        self._event.clear()
        # 每个线程都执行 _run 方法
        t = threading.Thread(target=self._run)
        t.start()

    def _run(self):
        '''
        开始执行线程
        :return: 
        '''
        # 获取当前线程的名字, 并加入到线程列表中
        current_thread_name = threading.currentThread().getName()
        self._all_thread_list.append(current_thread_name)

        # 开始任务循环
        # 等待获得任务
        task = self._task_queue.get()
        while True:

            # 如果是停止任务, 则跳出整个循环
            if self.STOP_TASK == task or self.TERMINATE == self._state:
                break

            if isinstance(task, Task):
                if task.is_ready():
                    # 开始执行任务
                    task.exec()

            # 等待获取任务
            with self._wait_queue(current_thread_name):
                if self.TERMINATE == self._state:
                    break

                # 获取任务
                task = self._task_queue.get()

        # 移除当前线程
        self._all_thread_list.remove(current_thread_name)
        if 0 == len(self._all_thread_list):
            self._event.set()

    # 请见 https://my.oschina.net/orangef/blog/133893
    @contextlib.contextmanager
    def _wait_queue(self, thread_name):
        """
        用于记录空闲的线程，或从空闲列表中取出线程处理任务
        """
        # 将当前线程名称添加到空闲线程列表中
        self._free_thread_list.append(thread_name)
        try:
            # 在此等待
            yield
        finally:
            # 将线程从空闲列表中移除
            self._free_thread_list.remove(thread_name)


def f1(i, value):
    time.sleep(2)
    current_thread_name = threading.currentThread().getName()
    print('%s: %s-%s' % (current_thread_name, i, value))
    if i == 3:
        return False
    return True


def test():
    pool = ThreadPool(5)
    for i in range(18):
        pool.apply_async(f1, args=(i, 'Test_%02d' % i,))
    pool.close()


if __name__ == '__main__':
    test()